import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet101_Weights

class DogClassifier(nn.Module):
    def __init__(self, num_classes=120):
        super(DogClassifier, self).__init__()
        # 使用ResNet101作为基础模型，使用最新的权重
        weights = ResNet101_Weights.IMAGENET1K_V2
        self.model = models.resnet101(weights=weights)
        
        # 修改分类器层
        num_features = self.model.fc.in_features
        self.model.fc = nn.Sequential(
            nn.Dropout(p=0.3),
            nn.Linear(num_features, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.3),
            nn.Linear(1024, num_classes)
        )
        
        # 初始化新层的权重
        for m in self.model.fc.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        
        # 默认冻结特征提取层
        self.freeze_features()
        
    def forward(self, x):
        return self.model(x)
    
    def freeze_features(self):
        """冻结特征提取层"""
        for name, param in self.model.named_parameters():
            if "fc" not in name:  # 不冻结全连接层
                param.requires_grad = False
            
    def unfreeze_features(self, num_layers=3):
        """解冻最后几个残差块进行微调"""
        # 首先冻结所有层
        self.freeze_features()
        
        # ResNet的层结构
        unfreeze_layers = [f'layer{i}' for i in range(5-num_layers, 5)]
        
        # 解冻指定的层
        for name, param in self.model.named_parameters():
            if any(layer in name for layer in unfreeze_layers) or "fc" in name:
                param.requires_grad = True
                
    def get_trainable_params(self):
        """返回需要训练的参数"""
        return [p for p in self.parameters() if p.requires_grad]